package com.yln.watcher.client;

import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import org.apache.zookeeper.*;
import org.apache.zookeeper.data.Stat;
import org.springframework.util.CollectionUtils;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Optional;

import static com.yln.watcher.client.PathConstant.*;
import static org.apache.zookeeper.ZooDefs.Ids.OPEN_ACL_UNSAFE;

@Slf4j
public class ZkClient implements Watcher {

    private ZooKeeper zooKeeper;
    private String hostPort;

    // 例：worker-1
    @Getter
    private String clientId;
    // 当前连接上zk server的客户端是否为master节点
    private boolean isMaster = false;

    // 客户端的上下文
    private ZkContext zkContext;

    public ZkClient(String hostPort, String clientId) {
        this.hostPort = hostPort;
        this.clientId = clientId;
    }

    // 连接上zk server
    public void start() throws IOException {
        zooKeeper = new ZooKeeper(hostPort, 15000, this);
    }

    // 主动断开连接
    public void stop() {
        try {
            if (zooKeeper != null) {
                zooKeeper.close();
            }
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
    }

    // 获取worker在workers节点下的全路径
    public String getWorkerNodePath() {
        // 只有非master节点才有此路径
        assert !isMaster;
        return WORKER + clientId;
    }

    // 获取/assign节点下的当前worker的路径
    // 例：/assign/worker-1
    public String getAssignedNodePath() {
        return ASSIGN + "/" + clientId;
    }

    // master用来获取某个worker节点在/assign下的路径
    // @param workerId 例：worker-1
    public String getAssignedNodePath(String workerId) {
        assert isMaster;
        return ASSIGN + "/" + workerId;
    }

    public String getTaskNodePath(String task) {
        assert isMaster;
        return TASKS + "/" + task;
    }

    // 初始化客户端，抢注/master，如果抢注失败，成为worker
    public void initContext() {
        tryBecomeMaster();
    }

    // 尝试创建/master节点
    public void tryBecomeMaster() {
        zooKeeper.create(MASTER, clientId.getBytes(), OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL, new BecomeMasterCallback(), this);
    }

    // 成为master节点
    public void becomeMaster() {
        isMaster = true;
        zkContext = new MasterContext();
        // 收集当前所有workers信息 和 当前所有正在工作的workers信息
        zooKeeper.create(WORKERS, null, OPEN_ACL_UNSAFE, CreateMode.PERSISTENT, new CreateWorkersNodeCallback(), this);
        zooKeeper.create(TASKS, null, OPEN_ACL_UNSAFE, CreateMode.PERSISTENT, new CreateTasksNodeCallback(), this);
        zooKeeper.create(ASSIGN, null, OPEN_ACL_UNSAFE, CreateMode.PERSISTENT, new CreateAssignNodeCallback(), this);
    }

    // 工具方法，为节点添加watcher
    public void addWatcher(String nodePath, Watcher watcher) {
        try {
            log.info("{} add watcher for node path: {}", clientId, nodePath);
            zooKeeper.addWatch(nodePath, watcher, AddWatchMode.PERSISTENT);
        } catch (InterruptedException | KeeperException e) {
            log.error("node path {"+nodePath+"} add watcher error", e);
        }
    }

    // 获取缓存的workers
    public List<String> getCacheWorkers() {
        assert isMaster;
        return zkContext.getAllWorkers();
    }

    // 获取缓存的tasks
    public List<String> getCacheTasks() {
        assert isMaster;
        return zkContext.getAllTasks();
    }

    // 获取缓存的assigned tasks
    public List<String> getCacheAssignedTasks() {
        assert !isMaster;
        return zkContext.getHasProceededTasks();
    }

    // 把/assign/{worker-id}下的所有task移动到/tasks下 等待重新分配
    public void rollbackTask(List<String> workers) {
        assert isMaster;
        if (CollectionUtils.isEmpty(workers)) {
            return;
        }
        workers.forEach(worker -> {
            try {
                List<String> tasks = zooKeeper.getChildren(getAssignedNodePath(worker), null);
                if (CollectionUtils.isEmpty(tasks)) {
                    return;
                }
                tasks.forEach(task -> {
                    String assignedTaskNodePath = getAssignedNodePath(worker) + "/" + task;
                    try {
                        Stat stat = new Stat();
                        byte[] data = zooKeeper.getData(assignedTaskNodePath, false, stat);
                        Op deleteAssignedTask = Op.delete(assignedTaskNodePath, stat.getVersion());
                        Op createTask = Op.create(getTaskNodePath(task), data, OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
                        zooKeeper.multi(Arrays.asList(deleteAssignedTask, createTask));
//                        zooKeeper.delete(assignedTaskNodePath, stat.getVersion());
//                        zooKeeper.create(getTaskNodePath(task), data, OPEN_ACL_UNSAFE, CreateMode.PERSISTENT, stat);
                    } catch (InterruptedException | KeeperException e) {
                        log.error("get assigned task {"+ assignedTaskNodePath + "/" + task +"} data error. ", e);
                    }
                });
            } catch (InterruptedException | KeeperException e) {
                log.error("rollback worker {"+ worker+"} tasks error. ", e);
            }
        });
    }

    // 成为worker节点
    public void becomeWorker() {
        isMaster = false;
        zkContext = new WorkerContext();
        registerWorkerNode();
    }

    public List<String> findAllWorkers() {
        return findChildrenNodes(WORKERS);
    }

    public List<String> findAllTasks() {
        return findChildrenNodes(TASKS);
    }

    public List<String> findAllAssignedTasksByClient() {
        return findChildrenNodes(getAssignedNodePath());
    }

    // 返回父节点下的所有子节点，返回值不包含全路径
    private List<String> findChildrenNodes(String parentNodeName) {
        List<String> children = new ArrayList<>();
        try {
            children = zooKeeper.getChildren(parentNodeName, false);
        } catch (InterruptedException | KeeperException e) {
            log.error("get children list error happened. ", e);
        }
        return children;
    }

    public void cacheWorkers(List<String> workers) {
        zkContext.cacheWorkers(workers);
    }

    public void cacheTasks(List<String> tasks) {
        zkContext.cacheTasks(tasks);
    }

    public void registerWorkerNode() {
        zooKeeper.create(getWorkerNodePath(), null, OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL, new RegisterWorkerCallback(), this);
    }

    public List<String> getAllAssignedTasks() {
        assert !isMaster;
        return getAllChildren(getAssignedNodePath());
    }

    private List<String> getAllChildren(String nodePath) {
        List<String> children = new ArrayList<>();
        try {
            children = zooKeeper.getChildren(nodePath, false);
        } catch (InterruptedException | KeeperException e) {
            log.error("get node {" + nodePath + "} children failed. ", e);
        }
        return children;
    }

    public String getNodeValue(String nodePath) {
        String payLoad = "";
        try {
            byte[] data = zooKeeper.getData(getAssignedNodePath() + "/" + nodePath, false, new Stat());
            return new String(data);
        } catch (InterruptedException | KeeperException e) {
            log.error("get node {" + nodePath + "} data error. ", e);
        }
        return payLoad;
    }

    public List<String> findAddedTasks(List<String> allAssignedTasks) {
        if (CollectionUtils.isEmpty(allAssignedTasks)) {
            return new ArrayList<>();
        }
        List<String> hasProceededTasks = zkContext.getHasProceededTasks();
        allAssignedTasks.retainAll(hasProceededTasks);
        return allAssignedTasks;
    }

    public void randomAssignTask(List<String> tasks) {
        log.info("--------- assign task to worker start ----------");
        assert isMaster;
        if (CollectionUtils.isEmpty(tasks)) {
            return;
        }
        tasks.forEach(task -> {
            Optional<String> workerOptional = zkContext.randomWorker();
            if (workerOptional.isPresent()) {
                String worker = workerOptional.get();
                log.info("get a worker {} to assign task.", worker);
                String assignTaskNodePath = getAssignedNodePath(worker) + "/" + task;
                log.info("assigned task node path: {}", assignTaskNodePath);
                String taskNodePath = getTaskNodePath(task);
                log.info("task node path: {}", taskNodePath);
                try {
                    Stat stat = new Stat();
                    byte[] data = zooKeeper.getData(taskNodePath, false, stat);
                    log.info("task {} data is {}", taskNodePath, new String(data));
                    Op deleteTask = Op.delete(taskNodePath, stat.getVersion());
                    Op createAssignTask = Op.create(assignTaskNodePath, data, OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
                    List<Op> opList = Arrays.asList(deleteTask, createAssignTask);
//                    zooKeeper.delete(taskNodePath, stat.getVersion());
//                    zooKeeper.create(assignTaskNodePath, data, OPEN_ACL_UNSAFE, CreateMode.PERSISTENT, new Stat());
                    zooKeeper.multi(opList);
                    zkContext.removeTask(task);
                } catch (InterruptedException | KeeperException e) {
                    log.error("assign task {"+assignTaskNodePath+"} error. ", e);
                }
            }
        });

    }

    public void createAssignWorkerNodeAndWatch() {
        // 非master才能创建自己的/assign/{worker-id}
        assert !isMaster;
        zooKeeper.create(getAssignedNodePath(), clientId.getBytes(), OPEN_ACL_UNSAFE, CreateMode.PERSISTENT, new CreateAssignWorkerCallback(), this);
    }

    @Override
    public void process(WatchedEvent event) {
        log.info("event path: {}, state: {}, type: {}", event.getPath(), event.getState(), event.getType());
    }

    public void doJob(String taskNode) {
        log.info("do job, taskNode is {}", taskNode);
        zkContext.taskProcessing(taskNode);
        String payLoad = getNodeValue(taskNode);
        doBusinessJob(payLoad);
    }

    private void doBusinessJob(String payLoad) {
        log.info("[Job started] payload: {}", payLoad);
        log.info("[Job ended] payload: {}", payLoad);
    }


}
